{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T07:49:36.630170Z",
     "start_time": "2021-10-11T07:49:33.253961Z"
    }
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.simplefilter('ignore')\n",
    "\n",
    "import os\n",
    "import time\n",
    "import copy\n",
    "import random\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "pd.set_option('max_columns', None)\n",
    "pd.set_option('max_colwidth', 400)\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import AutoTokenizer, AutoModel, get_linear_schedule_with_warmup, AutoConfig\n",
    "\n",
    "# 参数设置\n",
    "\n",
    "class Config:\n",
    "    def __init__(self):\n",
    "        super(Config, self).__init__()\n",
    "\n",
    "        self.SEED = 71\n",
    "        self.MODEL_PATH = 'hfl/chinese-roberta-wwm-ext'\n",
    "        self.NUM_LABELS = 6\n",
    "\n",
    "        # data\n",
    "        self.TOKENIZER = AutoTokenizer.from_pretrained(self.MODEL_PATH)\n",
    "        self.MAX_LENGTH = 400\n",
    "        self.BATCH_SIZE = 8\n",
    "        self.VALIDATION_SPLIT = 0.10\n",
    "\n",
    "        # model\n",
    "        self.DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "        self.FULL_FINETUNING = True\n",
    "        self.LR = 3e-5\n",
    "        self.OPTIMIZER = 'AdamW'\n",
    "        self.N_VALIDATE_DUR_TRAIN = 3\n",
    "        self.N_WARMUP = 0\n",
    "        self.SAVE_BEST_ONLY = True\n",
    "        self.EPOCHS = 1\n",
    "        self.USE_FGM = False\n",
    "\n",
    "\n",
    "config = Config()\n",
    "\n",
    "\n",
    "def seed_torch(seed=42):\n",
    "    random.seed(seed)\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "\n",
    "\n",
    "np.random.seed(config.SEED)\n",
    "seed_torch(seed=config.SEED)\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T07:49:38.258719Z",
     "start_time": "2021-10-11T07:49:37.708426Z"
    }
   },
   "outputs": [],
   "source": [
    "train = pd.read_csv('data/train_scenes_characters.csv')\n",
    "test = pd.read_csv('data/test_scenes_characters.csv')\n",
    "\n",
    "submit = pd.read_csv('raw_data/submit_example.tsv', sep='\\t')\n",
    "\n",
    "train['labels'] = train['emotions'].apply(lambda x: [int(i) for i in x.split(',')])\n",
    "for i in range(6):\n",
    "    train[f'label_{i}'] = train['labels'].apply(lambda x: x[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T07:49:39.171922Z",
     "start_time": "2021-10-11T07:49:39.147043Z"
    }
   },
   "outputs": [],
   "source": [
    "train_df = train[['id', 'movie', 'text'] + [f'label_{i}' for i in range(6)]].copy().reset_index(drop=True)\n",
    "test_df = test[['id', 'movie', 'text']].copy().reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T07:49:40.454940Z",
     "start_time": "2021-10-11T07:49:40.422738Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(36782, 9)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>movie</th>\n",
       "      <th>text</th>\n",
       "      <th>label_0</th>\n",
       "      <th>label_1</th>\n",
       "      <th>label_2</th>\n",
       "      <th>label_3</th>\n",
       "      <th>label_4</th>\n",
       "      <th>label_5</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1171_0001_A_1</td>\n",
       "      <td>1171</td>\n",
       "      <td>剧本: 1171 场景: 1 当前: 天空下着暴雨， 剧本1171 何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。 角色: 何仁晴 前文:</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1171_0001_A_3</td>\n",
       "      <td>1171</td>\n",
       "      <td>剧本: 1171 场景: 1 当前:  剧本1171 何仁晴一手拿着一个行李，一路小跑着把刘昆诚带到了文工团门口。 角色: 何仁晴 前文: 天空下着暴雨，何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1171_0001_A_5</td>\n",
       "      <td>1171</td>\n",
       "      <td>剧本: 1171 场景: 1 当前:  剧本1171 何仁晴停下来接过刘昆诚手里的行李：你妈妈交待我了，等领了军装一定要照张相寄回去，让街坊邻居都知道你当兵了。 角色: 何仁晴 前文: 何仁晴一手拿着一个行李，一路小跑着把刘昆诚带到了文工团门口。天空下着暴雨，何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1171_0001_A_8</td>\n",
       "      <td>1171</td>\n",
       "      <td>剧本: 1171 场景: 1 当前:  剧本1171 何仁晴凑近刘昆诚小声：办入伍证审的时候，派出所的民警跟我说，你的亲生父亲还在劳改，但是你跟他划清了界限，改姓了你继父的姓，所以出身这一栏，我就给你填革干了，进了团不要跟别人说这件事，我也不会说的。 角色: 何仁晴 前文: 何仁晴停下来接过刘昆诚手里的行李：你妈妈交待我了，等领了军装一定要照张相寄回去，让街坊邻居都知道你当兵了。何仁晴一手拿着一个行李，一路小跑着把刘昆诚带到了文工团门口。天空下着暴雨，何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1171_0001_A_11</td>\n",
       "      <td>1171</td>\n",
       "      <td>剧本: 1171 场景: 1 当前:  剧本1171 何仁晴笑了笑：军礼不是这么敬的。五指并拢，大臂带动小臂，举到齐眉。 角色: 何仁晴 前文: 何仁晴凑近刘昆诚小声：办入伍证审的时候，派出所的民警跟我说，你的亲生父亲还在劳改，但是你跟他划清了界限，改姓了你继父的姓，所以出身这一栏，我就给你填革干了，进了团不要跟别人说这件事，我也不会说的。何仁晴停下来接过刘昆诚手里的行李：你妈妈交待我了，等领了军装一定要照张相寄回去，让街坊邻居都知道你当兵了。何仁晴一手拿着一个行李，一路小跑着把刘昆诚带到了文工团门口。天空下着暴雨，何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>1171_0001_A_12</td>\n",
       "      <td>1171</td>\n",
       "      <td>剧本: 1171 场景: 1 当前:  剧本1171 何仁晴示范了一个动作，刘昆诚照做。 角色: 何仁晴 前文: 何仁晴笑了笑：军礼不是这么敬的。五指并拢，大臂带动小臂，举到齐眉。何仁晴凑近刘昆诚小声：办入伍证审的时候，派出所的民警跟我说，你的亲生父亲还在劳改，但是你跟他划清了界限，改姓了你继父的姓，所以出身这一栏，我就给你填革干了，进了团不要跟别人说这件事，我也不会说的。何仁晴停下来接过刘昆诚手里的行李：你妈妈交待我了，等领了军装一定要照张相寄回去，让街坊邻居都知道你当兵了。何仁晴一手拿着一个行李，一路小跑着把刘昆诚带到了文工团门口。天空下着暴雨，何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>1171_0001_A_14</td>\n",
       "      <td>1171</td>\n",
       "      <td>剧本: 1171 场景: 1 当前:  剧本1171 何仁晴：礼毕。（再次举手敬礼）敬礼。 角色: 何仁晴 前文: 何仁晴示范了一个动作，刘昆诚照做。何仁晴笑了笑：军礼不是这么敬的。五指并拢，大臂带动小臂，举到齐眉。何仁晴凑近刘昆诚小声：办入伍证审的时候，派出所的民警跟我说，你的亲生父亲还在劳改，但是你跟他划清了界限，改姓了你继父的姓，所以出身这一栏，我就给你填革干了，进了团不要跟别人说这件事，我也不会说的。何仁晴停下来接过刘昆诚手里的行李：你妈妈交待我了，等领了军装一定要照张相寄回去，让街坊邻居都知道你当兵了。何仁晴一手拿着一个行李，一路小跑着把刘昆诚带到了文工团门口。天空下着暴雨，何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>1171_0001_A_2</td>\n",
       "      <td>1171</td>\n",
       "      <td>剧本: 1171 场景: 1 当前: 天空下着暴雨，何仁晴正在给 剧本1171 刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。 角色: 刘昆诚 前文:</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>1171_0001_A_4</td>\n",
       "      <td>1171</td>\n",
       "      <td>剧本: 1171 场景: 1 当前: 何仁晴一手拿着一个行李，一路小跑着把 剧本1171 刘昆诚带到了文工团门口。 角色: 刘昆诚 前文: 天空下着暴雨，何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>1171_0001_A_6</td>\n",
       "      <td>1171</td>\n",
       "      <td>剧本: 1171 场景: 1 当前: 何仁晴停下来接过 剧本1171 刘昆诚手里的行李：你妈妈交待我了，等领了军装一定要照张相寄回去，让街坊邻居都知道你当兵了。 角色: 刘昆诚 前文: 何仁晴一手拿着一个行李，一路小跑着把刘昆诚带到了文工团门口。天空下着暴雨，何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               id  movie  \\\n",
       "0   1171_0001_A_1   1171   \n",
       "1   1171_0001_A_3   1171   \n",
       "2   1171_0001_A_5   1171   \n",
       "3   1171_0001_A_8   1171   \n",
       "4  1171_0001_A_11   1171   \n",
       "5  1171_0001_A_12   1171   \n",
       "6  1171_0001_A_14   1171   \n",
       "7   1171_0001_A_2   1171   \n",
       "8   1171_0001_A_4   1171   \n",
       "9   1171_0001_A_6   1171   \n",
       "\n",
       "                                                                                                                                                                                                                                                                                                                                           text  \\\n",
       "0                                                                                                                                                                                                                                                           剧本: 1171 场景: 1 当前: 天空下着暴雨， 剧本1171 何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。 角色: 何仁晴 前文:    \n",
       "1                                                                                                                                                                                                                             剧本: 1171 场景: 1 当前:  剧本1171 何仁晴一手拿着一个行李，一路小跑着把刘昆诚带到了文工团门口。 角色: 何仁晴 前文: 天空下着暴雨，何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。   \n",
       "2                                                                                                                                                                        剧本: 1171 场景: 1 当前:  剧本1171 何仁晴停下来接过刘昆诚手里的行李：你妈妈交待我了，等领了军装一定要照张相寄回去，让街坊邻居都知道你当兵了。 角色: 何仁晴 前文: 何仁晴一手拿着一个行李，一路小跑着把刘昆诚带到了文工团门口。天空下着暴雨，何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。   \n",
       "3                                                                      剧本: 1171 场景: 1 当前:  剧本1171 何仁晴凑近刘昆诚小声：办入伍证审的时候，派出所的民警跟我说，你的亲生父亲还在劳改，但是你跟他划清了界限，改姓了你继父的姓，所以出身这一栏，我就给你填革干了，进了团不要跟别人说这件事，我也不会说的。 角色: 何仁晴 前文: 何仁晴停下来接过刘昆诚手里的行李：你妈妈交待我了，等领了军装一定要照张相寄回去，让街坊邻居都知道你当兵了。何仁晴一手拿着一个行李，一路小跑着把刘昆诚带到了文工团门口。天空下着暴雨，何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。   \n",
       "4                                     剧本: 1171 场景: 1 当前:  剧本1171 何仁晴笑了笑：军礼不是这么敬的。五指并拢，大臂带动小臂，举到齐眉。 角色: 何仁晴 前文: 何仁晴凑近刘昆诚小声：办入伍证审的时候，派出所的民警跟我说，你的亲生父亲还在劳改，但是你跟他划清了界限，改姓了你继父的姓，所以出身这一栏，我就给你填革干了，进了团不要跟别人说这件事，我也不会说的。何仁晴停下来接过刘昆诚手里的行李：你妈妈交待我了，等领了军装一定要照张相寄回去，让街坊邻居都知道你当兵了。何仁晴一手拿着一个行李，一路小跑着把刘昆诚带到了文工团门口。天空下着暴雨，何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。   \n",
       "5                    剧本: 1171 场景: 1 当前:  剧本1171 何仁晴示范了一个动作，刘昆诚照做。 角色: 何仁晴 前文: 何仁晴笑了笑：军礼不是这么敬的。五指并拢，大臂带动小臂，举到齐眉。何仁晴凑近刘昆诚小声：办入伍证审的时候，派出所的民警跟我说，你的亲生父亲还在劳改，但是你跟他划清了界限，改姓了你继父的姓，所以出身这一栏，我就给你填革干了，进了团不要跟别人说这件事，我也不会说的。何仁晴停下来接过刘昆诚手里的行李：你妈妈交待我了，等领了军装一定要照张相寄回去，让街坊邻居都知道你当兵了。何仁晴一手拿着一个行李，一路小跑着把刘昆诚带到了文工团门口。天空下着暴雨，何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。   \n",
       "6  剧本: 1171 场景: 1 当前:  剧本1171 何仁晴：礼毕。（再次举手敬礼）敬礼。 角色: 何仁晴 前文: 何仁晴示范了一个动作，刘昆诚照做。何仁晴笑了笑：军礼不是这么敬的。五指并拢，大臂带动小臂，举到齐眉。何仁晴凑近刘昆诚小声：办入伍证审的时候，派出所的民警跟我说，你的亲生父亲还在劳改，但是你跟他划清了界限，改姓了你继父的姓，所以出身这一栏，我就给你填革干了，进了团不要跟别人说这件事，我也不会说的。何仁晴停下来接过刘昆诚手里的行李：你妈妈交待我了，等领了军装一定要照张相寄回去，让街坊邻居都知道你当兵了。何仁晴一手拿着一个行李，一路小跑着把刘昆诚带到了文工团门口。天空下着暴雨，何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。   \n",
       "7                                                                                                                                                                                                                                                           剧本: 1171 场景: 1 当前: 天空下着暴雨，何仁晴正在给 剧本1171 刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。 角色: 刘昆诚 前文:    \n",
       "8                                                                                                                                                                                                                             剧本: 1171 场景: 1 当前: 何仁晴一手拿着一个行李，一路小跑着把 剧本1171 刘昆诚带到了文工团门口。 角色: 刘昆诚 前文: 天空下着暴雨，何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。   \n",
       "9                                                                                                                                                                        剧本: 1171 场景: 1 当前: 何仁晴停下来接过 剧本1171 刘昆诚手里的行李：你妈妈交待我了，等领了军装一定要照张相寄回去，让街坊邻居都知道你当兵了。 角色: 刘昆诚 前文: 何仁晴一手拿着一个行李，一路小跑着把刘昆诚带到了文工团门口。天空下着暴雨，何仁晴正在给刘昆诚穿雨衣，他自己却只穿着单薄的军装，完全暴露在大雨之中。   \n",
       "\n",
       "   label_0  label_1  label_2  label_3  label_4  label_5  \n",
       "0        0        0        0        0        0        0  \n",
       "1        0        0        0        0        0        0  \n",
       "2        0        0        0        0        0        0  \n",
       "3        0        0        0        0        0        0  \n",
       "4        0        1        0        0        0        0  \n",
       "5        0        0        0        0        0        0  \n",
       "6        0        0        0        0        0        0  \n",
       "7        0        0        0        0        0        0  \n",
       "8        0        0        0        0        0        0  \n",
       "9        0        0        0        0        0        0  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(train_df.shape)\n",
    "train_df.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T07:49:42.247033Z",
     "start_time": "2021-10-11T07:49:42.226449Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(21376, 3)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>movie</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1597_0001_A_7</td>\n",
       "      <td>1597</td>\n",
       "      <td>剧本: 1597 场景: 1 当前: 宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子 剧本1597 周惠骏，偌大的会议室只有他们俩。 角色: 周惠骏 前文:</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1597_0001_A_9</td>\n",
       "      <td>1597</td>\n",
       "      <td>剧本: 1597 场景: 1 当前:  剧本1597 周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？ 角色: 周惠骏 前文: 宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1597_0001_A_11</td>\n",
       "      <td>1597</td>\n",
       "      <td>剧本: 1597 场景: 1 当前:  剧本1597 周惠骏：罗彤伯出事了。 角色: 周惠骏 前文: 周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1597_0001_A_13</td>\n",
       "      <td>1597</td>\n",
       "      <td>剧本: 1597 场景: 1 当前:  剧本1597 周惠骏：她为了证明他去过海洋馆，昨天跟蛋蛋打了一架，蛋蛋推了她一下，把她头给磕破了。 角色: 周惠骏 前文: 周惠骏：罗彤伯出事了。周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1597_0001_A_15</td>\n",
       "      <td>1597</td>\n",
       "      <td>剧本: 1597 场景: 1 当前: 宋淳云有些生气走向 剧本1597 周惠骏：怎么会这样呢？我说过我会带她去的，你为什么不看好她呢？ 角色: 周惠骏 前文: 周惠骏：她为了证明他去过海洋馆，昨天跟蛋蛋打了一架，蛋蛋推了她一下，把她头给磕破了。周惠骏：罗彤伯出事了。周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>1597_0001_A_16</td>\n",
       "      <td>1597</td>\n",
       "      <td>剧本: 1597 场景: 1 当前:  剧本1597 周惠骏：你已经推了七次了，她一直都在等你，你知道吗？ 角色: 周惠骏 前文: 宋淳云有些生气走向周惠骏：怎么会这样呢？我说过我会带她去的，你为什么不看好她呢？周惠骏：她为了证明他去过海洋馆，昨天跟蛋蛋打了一架，蛋蛋推了她一下，把她头给磕破了。周惠骏：罗彤伯出事了。周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>1597_0001_A_18</td>\n",
       "      <td>1597</td>\n",
       "      <td>剧本: 1597 场景: 1 当前:  剧本1597 周惠骏态度坚决：我只要罗彤伯。 角色: 周惠骏 前文: 周惠骏：你已经推了七次了，她一直都在等你，你知道吗？宋淳云有些生气走向周惠骏：怎么会这样呢？我说过我会带她去的，你为什么不看好她呢？周惠骏：她为了证明他去过海洋馆，昨天跟蛋蛋打了一架，蛋蛋推了她一下，把她头给磕破了。周惠骏：罗彤伯出事了。周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>1597_0001_A_22</td>\n",
       "      <td>1597</td>\n",
       "      <td>剧本: 1597 场景: 1 当前:  剧本1597 周惠骏一笑：上次你下跪，是向我求婚，真幽默。高苹霞敲了敲办公室的玻璃门，他脖子落枕，有些不舒服。 角色: 周惠骏 前文: 周惠骏态度坚决：我只要罗彤伯。周惠骏：你已经推了七次了，她一直都在等你，你知道吗？宋淳云有些生气走向周惠骏：怎么会这样呢？我说过我会带她去的，你为什么不看好她呢？周惠骏：她为了证明他去过海洋馆，昨天跟蛋蛋打了一架，蛋蛋推了她一下，把她头给磕破了。周惠骏：罗彤伯出事了。周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>1597_0001_A_24</td>\n",
       "      <td>1597</td>\n",
       "      <td>剧本: 1597 场景: 1 当前:  剧本1597 周惠骏将头扭向一侧，宋淳云转头望了望。 角色: 周惠骏 前文: 周惠骏一笑：上次你下跪，是向我求婚，真幽默。高苹霞敲了敲办公室的玻璃门，他脖子落枕，有些不舒服。周惠骏态度坚决：我只要罗彤伯。周惠骏：你已经推了七次了，她一直都在等你，你知道吗？宋淳云有些生气走向周惠骏：怎么会这样呢？我说过我会带她去的，你为什么不看好她呢？周惠骏：她为了证明他去过海洋馆，昨天跟蛋蛋打了一架，蛋蛋推了她一下，把她头给磕破了。周惠骏：罗彤伯出事了。周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>1597_0001_A_42</td>\n",
       "      <td>1597</td>\n",
       "      <td>剧本: 1597 场景: 1 当前: 宋淳云拦住她， 剧本1597 周惠骏挣脱：放开。你刚才不是说要去T国吗？ 角色: 周惠骏 前文: 周惠骏将头扭向一侧，宋淳云转头望了望。周惠骏一笑：上次你下跪，是向我求婚，真幽默。高苹霞敲了敲办公室的玻璃门，他脖子落枕，有些不舒服。周惠骏态度坚决：我只要罗彤伯。周惠骏：你已经推了七次了，她一直都在等你，你知道吗？宋淳云有些生气走向周惠骏：怎么会这样呢？我说过我会带她去的，你为什么不看好她呢？周惠骏：她为了证明他去过海洋馆，昨天跟蛋蛋打了一架，蛋蛋推了她一下，把她头给磕破了。周惠骏：罗彤伯出事了。周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               id  movie  \\\n",
       "0   1597_0001_A_7   1597   \n",
       "1   1597_0001_A_9   1597   \n",
       "2  1597_0001_A_11   1597   \n",
       "3  1597_0001_A_13   1597   \n",
       "4  1597_0001_A_15   1597   \n",
       "5  1597_0001_A_16   1597   \n",
       "6  1597_0001_A_18   1597   \n",
       "7  1597_0001_A_22   1597   \n",
       "8  1597_0001_A_24   1597   \n",
       "9  1597_0001_A_42   1597   \n",
       "\n",
       "                                                                                                                                                                                                                                                                                                                                                                                                  text  \n",
       "0                                                                                                                                                                                                                                                                                                              剧本: 1597 场景: 1 当前: 宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子 剧本1597 周惠骏，偌大的会议室只有他们俩。 角色: 周惠骏 前文:   \n",
       "1                                                                                                                                                                                                                                       剧本: 1597 场景: 1 当前:  剧本1597 周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？ 角色: 周惠骏 前文: 宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。  \n",
       "2                                                                                                                                                                                                                            剧本: 1597 场景: 1 当前:  剧本1597 周惠骏：罗彤伯出事了。 角色: 周惠骏 前文: 周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。  \n",
       "3                                                                                                                                                                                  剧本: 1597 场景: 1 当前:  剧本1597 周惠骏：她为了证明他去过海洋馆，昨天跟蛋蛋打了一架，蛋蛋推了她一下，把她头给磕破了。 角色: 周惠骏 前文: 周惠骏：罗彤伯出事了。周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。  \n",
       "4                                                                                                                                          剧本: 1597 场景: 1 当前: 宋淳云有些生气走向 剧本1597 周惠骏：怎么会这样呢？我说过我会带她去的，你为什么不看好她呢？ 角色: 周惠骏 前文: 周惠骏：她为了证明他去过海洋馆，昨天跟蛋蛋打了一架，蛋蛋推了她一下，把她头给磕破了。周惠骏：罗彤伯出事了。周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。  \n",
       "5                                                                                                                剧本: 1597 场景: 1 当前:  剧本1597 周惠骏：你已经推了七次了，她一直都在等你，你知道吗？ 角色: 周惠骏 前文: 宋淳云有些生气走向周惠骏：怎么会这样呢？我说过我会带她去的，你为什么不看好她呢？周惠骏：她为了证明他去过海洋馆，昨天跟蛋蛋打了一架，蛋蛋推了她一下，把她头给磕破了。周惠骏：罗彤伯出事了。周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。  \n",
       "6                                                                                                 剧本: 1597 场景: 1 当前:  剧本1597 周惠骏态度坚决：我只要罗彤伯。 角色: 周惠骏 前文: 周惠骏：你已经推了七次了，她一直都在等你，你知道吗？宋淳云有些生气走向周惠骏：怎么会这样呢？我说过我会带她去的，你为什么不看好她呢？周惠骏：她为了证明他去过海洋馆，昨天跟蛋蛋打了一架，蛋蛋推了她一下，把她头给磕破了。周惠骏：罗彤伯出事了。周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。  \n",
       "7                                                 剧本: 1597 场景: 1 当前:  剧本1597 周惠骏一笑：上次你下跪，是向我求婚，真幽默。高苹霞敲了敲办公室的玻璃门，他脖子落枕，有些不舒服。 角色: 周惠骏 前文: 周惠骏态度坚决：我只要罗彤伯。周惠骏：你已经推了七次了，她一直都在等你，你知道吗？宋淳云有些生气走向周惠骏：怎么会这样呢？我说过我会带她去的，你为什么不看好她呢？周惠骏：她为了证明他去过海洋馆，昨天跟蛋蛋打了一架，蛋蛋推了她一下，把她头给磕破了。周惠骏：罗彤伯出事了。周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。  \n",
       "8                              剧本: 1597 场景: 1 当前:  剧本1597 周惠骏将头扭向一侧，宋淳云转头望了望。 角色: 周惠骏 前文: 周惠骏一笑：上次你下跪，是向我求婚，真幽默。高苹霞敲了敲办公室的玻璃门，他脖子落枕，有些不舒服。周惠骏态度坚决：我只要罗彤伯。周惠骏：你已经推了七次了，她一直都在等你，你知道吗？宋淳云有些生气走向周惠骏：怎么会这样呢？我说过我会带她去的，你为什么不看好她呢？周惠骏：她为了证明他去过海洋馆，昨天跟蛋蛋打了一架，蛋蛋推了她一下，把她头给磕破了。周惠骏：罗彤伯出事了。周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。  \n",
       "9  剧本: 1597 场景: 1 当前: 宋淳云拦住她， 剧本1597 周惠骏挣脱：放开。你刚才不是说要去T国吗？ 角色: 周惠骏 前文: 周惠骏将头扭向一侧，宋淳云转头望了望。周惠骏一笑：上次你下跪，是向我求婚，真幽默。高苹霞敲了敲办公室的玻璃门，他脖子落枕，有些不舒服。周惠骏态度坚决：我只要罗彤伯。周惠骏：你已经推了七次了，她一直都在等你，你知道吗？宋淳云有些生气走向周惠骏：怎么会这样呢？我说过我会带她去的，你为什么不看好她呢？周惠骏：她为了证明他去过海洋馆，昨天跟蛋蛋打了一架，蛋蛋推了她一下，把她头给磕破了。周惠骏：罗彤伯出事了。周惠骏有些不高兴：宋淳云，我们俩在一起十二年了，在我跟你提出离婚的时候，你第一个反应居然是不要影响你的生意，这场婚姻对于你来说毫无意义，是吗？宋淳云将视频暂停，转身面对会议桌，桌子最那头，坐着宋淳云的妻子周惠骏，偌大的会议室只有他们俩。  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(test_df.shape)\n",
    "test_df.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T07:50:02.722160Z",
     "start_time": "2021-10-11T07:50:02.701984Z"
    }
   },
   "outputs": [],
   "source": [
    "class TransformerDataset(Dataset):\n",
    "    def __init__(self, df, indices, set_type=None):\n",
    "        super(TransformerDataset, self).__init__()\n",
    "\n",
    "        df = df.iloc[indices]\n",
    "        self.texts = df['text'].values.tolist()\n",
    "        self.set_type = set_type\n",
    "        if self.set_type != 'test':\n",
    "            self.labels = df.iloc[:, 3:].values\n",
    "\n",
    "        self.tokenizer = config.TOKENIZER\n",
    "        self.max_length = config.MAX_LENGTH\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.texts)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        tokenized = self.tokenizer.encode_plus(\n",
    "            self.texts[index],\n",
    "            add_special_tokens=True,\n",
    "            max_length=self.max_length,\n",
    "            pad_to_max_length=True,\n",
    "            truncation=True,\n",
    "            return_attention_mask=True,\n",
    "            return_token_type_ids=False,\n",
    "            return_tensors='pt'\n",
    "        )\n",
    "        input_ids = tokenized['input_ids'].squeeze()\n",
    "        attention_mask = tokenized['attention_mask'].squeeze()\n",
    "\n",
    "        if self.set_type != 'test':\n",
    "            return {\n",
    "                'input_ids': input_ids.long(),\n",
    "                'attention_mask': attention_mask.long(),\n",
    "                'labels': torch.Tensor(self.labels[index]).float(),\n",
    "            }\n",
    "\n",
    "        return {\n",
    "            'input_ids': input_ids.long(),\n",
    "            'attention_mask': attention_mask.long(),\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T07:50:03.968143Z",
     "start_time": "2021-10-11T07:50:03.925834Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_ids shape: torch.Size([8, 400])\n",
      "attention_mask shape: torch.Size([8, 400])\n",
      "labels shape: torch.Size([8, 6])\n"
     ]
    }
   ],
   "source": [
    "dataset_size = len(train_df)\n",
    "indices = list(range(dataset_size))\n",
    "split = int(np.floor(config.VALIDATION_SPLIT * dataset_size))\n",
    "np.random.shuffle(indices)\n",
    "\n",
    "train_indices, valid_indices = indices[split:], indices[:split]\n",
    "\n",
    "train_data = TransformerDataset(train_df, train_indices)\n",
    "valid_data = TransformerDataset(train_df, valid_indices)\n",
    "\n",
    "train_dataloader = DataLoader(train_data, batch_size=config.BATCH_SIZE)\n",
    "valid_dataloader = DataLoader(valid_data, batch_size=config.BATCH_SIZE)\n",
    "\n",
    "b = next(iter(train_dataloader))\n",
    "for k, v in b.items():\n",
    "    print(f'{k} shape: {v.shape}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T07:50:11.884325Z",
     "start_time": "2021-10-11T07:50:11.867961Z"
    }
   },
   "outputs": [],
   "source": [
    "class FGM(object):\n",
    "    def __init__(self, model, emb_name, epsilon=1.0):\n",
    "        self.model = model\n",
    "        self.epsilon = epsilon\n",
    "        self.emb_name = emb_name\n",
    "        self.backup = {}\n",
    "\n",
    "    def attack(self):\n",
    "        for name, param in self.model.named_parameters():\n",
    "            if param.requires_grad and self.emb_name in name:\n",
    "                self.backup[name] = param.data.clone()\n",
    "                norm = torch.norm(param.grad)\n",
    "                if norm != 0 and not torch.isnan(norm):\n",
    "                    r_at = self.epsilon * param.grad / norm\n",
    "                    param.data.add_(r_at)\n",
    "\n",
    "    def restore(self):\n",
    "        for name, param in self.model.named_parameters():\n",
    "            if param.requires_grad and self.emb_name in name:\n",
    "                assert name in self.backup\n",
    "                param.data = self.backup[name]\n",
    "        self.backup = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T07:50:14.681613Z",
     "start_time": "2021-10-11T07:50:14.656922Z"
    }
   },
   "outputs": [],
   "source": [
    "def init_params(module_lst):\n",
    "    for module in module_lst:\n",
    "        for param in module.parameters():\n",
    "            if param.dim() > 1:\n",
    "                torch.nn.init.xavier_uniform_(param)\n",
    "    return\n",
    "\n",
    "\n",
    "class Model(nn.Module):\n",
    "    def __init__(self, ):\n",
    "        super().__init__()\n",
    "\n",
    "        cfg = AutoConfig.from_pretrained(config.MODEL_PATH)\n",
    "        cfg.update({\"output_hidden_states\": True,\n",
    "                    \"hidden_dropout_prob\": 0.0,\n",
    "                    \"layer_norm_eps\": 1e-7})\n",
    "\n",
    "        self.roberta = AutoModel.from_pretrained(config.MODEL_PATH, config=cfg)\n",
    "\n",
    "        dim = self.roberta.pooler.dense.bias.shape[0]\n",
    "\n",
    "        self.dropout = nn.Dropout(p=0.2)\n",
    "        self.high_dropout = nn.Dropout(p=0.5)\n",
    "\n",
    "        n_weights = 12\n",
    "        weights_init = torch.zeros(n_weights).float()\n",
    "        weights_init.data[:-1] = -3\n",
    "        self.layer_weights = torch.nn.Parameter(weights_init)\n",
    "\n",
    "        self.attention = nn.Sequential(\n",
    "            nn.Linear(768, 768),\n",
    "            nn.Tanh(),\n",
    "            nn.Linear(768, 1),\n",
    "            nn.Softmax(dim=1)\n",
    "        )\n",
    "        self.cls = nn.Sequential(\n",
    "            nn.Linear(dim, 6)\n",
    "        )\n",
    "        init_params([self.cls, self.attention])\n",
    "\n",
    "    def forward(self, input_ids, attention_mask):\n",
    "        roberta_output = self.roberta(input_ids=input_ids,\n",
    "                                      attention_mask=attention_mask)\n",
    "\n",
    "        cls_outputs = torch.stack(\n",
    "            [self.dropout(layer) for layer in roberta_output[2][-12:]], dim=0\n",
    "        )\n",
    "        cls_output = (\n",
    "                torch.softmax(self.layer_weights, dim=0).unsqueeze(1).unsqueeze(1).unsqueeze(1) * cls_outputs).sum(\n",
    "            0)\n",
    "\n",
    "        logits = torch.mean(\n",
    "            torch.stack(\n",
    "                [torch.sum(self.attention(self.high_dropout(cls_output)) * cls_output, dim=1) for _ in range(5)],\n",
    "                dim=0,\n",
    "            ),\n",
    "            dim=0,\n",
    "        )\n",
    "        return self.cls(logits)\n",
    "    \n",
    "device = config.DEVICE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T07:50:18.753737Z",
     "start_time": "2021-10-11T07:50:18.728596Z"
    }
   },
   "outputs": [],
   "source": [
    "def val(model, valid_dataloader, criterion):\n",
    "    val_loss = 0\n",
    "    true, pred = [], []\n",
    "\n",
    "    # set model.eval() every time during evaluation\n",
    "    model.eval()\n",
    "\n",
    "    for step, batch in enumerate(valid_dataloader):\n",
    "        b_input_ids = batch['input_ids'].to(device)\n",
    "        b_attention_mask = batch['attention_mask'].to(device)\n",
    "        b_labels = batch['labels'].to(device)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            # forward pass\n",
    "            logits = model(input_ids=b_input_ids, attention_mask=b_attention_mask)\n",
    "\n",
    "            # calculate loss\n",
    "            loss = criterion(logits, b_labels)\n",
    "            val_loss += loss.item()\n",
    "\n",
    "    avg_val_loss = val_loss / len(valid_dataloader)\n",
    "    print('Val loss:', avg_val_loss)\n",
    "    return avg_val_loss\n",
    "\n",
    "\n",
    "def train(model, train_dataloader, valid_dataloader, criterion, optimizer, scheduler, epoch):\n",
    "    # we validate config.N_VALIDATE_DUR_TRAIN times during the training loop\n",
    "    nv = config.N_VALIDATE_DUR_TRAIN\n",
    "    temp = len(train_dataloader) // nv\n",
    "    temp = temp - (temp % 100)\n",
    "    validate_at_steps = [temp * x for x in range(1, nv + 1)]\n",
    "    \n",
    "    if config.USE_FGM:\n",
    "        fgm = FGM(model, epsilon=1, emb_name='word_embeddings.')\n",
    "\n",
    "    train_loss = 0\n",
    "    for step, batch in enumerate(tqdm(train_dataloader,\n",
    "                                      desc='Epoch ' + str(epoch))):\n",
    "        # set model.eval() every time during training\n",
    "        model.train()\n",
    "\n",
    "        # unpack the batch contents and push them to the device (cuda or cpu).\n",
    "        b_input_ids = batch['input_ids'].to(device)\n",
    "        b_attention_mask = batch['attention_mask'].to(device)\n",
    "        b_labels = batch['labels'].to(device)\n",
    "\n",
    "        # clear accumulated gradients\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # forward pass\n",
    "        logits = model(input_ids=b_input_ids, attention_mask=b_attention_mask)\n",
    "\n",
    "        # calculate loss\n",
    "        loss = criterion(logits, b_labels)\n",
    "        train_loss += loss.item()\n",
    "\n",
    "        # backward pass\n",
    "        loss.backward()\n",
    "        \n",
    "        # fgm attack\n",
    "        if config.USE_FGM:\n",
    "            fgm.attack()\n",
    "            logits_adv = model(input_ids=b_input_ids, attention_mask=b_attention_mask)\n",
    "            loss_adv = criterion(logits_adv, b_labels)\n",
    "            loss_adv.backward()\n",
    "            fgm.restore()\n",
    "\n",
    "        # update weights\n",
    "        optimizer.step()\n",
    "\n",
    "        # update scheduler\n",
    "        scheduler.step()\n",
    "\n",
    "        if step in validate_at_steps:\n",
    "            print(f'-- Step: {step}')\n",
    "            _ = val(model, valid_dataloader, criterion)\n",
    "\n",
    "    avg_train_loss = train_loss / len(train_dataloader)\n",
    "    print('Training loss:', avg_train_loss)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T07:50:30.214551Z",
     "start_time": "2021-10-11T07:50:30.194210Z"
    }
   },
   "outputs": [],
   "source": [
    "class RMSELoss(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(RMSELoss,self).__init__()\n",
    "\n",
    "    def forward(self, x, y):\n",
    "        criterion = nn.MSELoss()\n",
    "        loss = torch.sqrt(criterion(x, y))\n",
    "        return loss\n",
    "    \n",
    "def run():\n",
    "    # setting a seed ensures reproducible results.\n",
    "    # seed may affect the performance too.\n",
    "    torch.manual_seed(config.SEED)\n",
    "\n",
    "    # criterion = nn.BCEWithLogitsLoss()\n",
    "    criterion = RMSELoss()\n",
    "\n",
    "    # define the parameters to be optmized -\n",
    "    # - and add regularization\n",
    "    if config.FULL_FINETUNING:\n",
    "        param_optimizer = list(model.named_parameters())\n",
    "        no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n",
    "        optimizer_parameters = [\n",
    "            {\n",
    "                \"params\": [\n",
    "                    p for n, p in param_optimizer if not any(nd in n for nd in no_decay)\n",
    "                ],\n",
    "                \"weight_decay\": 0.001,\n",
    "            },\n",
    "            {\n",
    "                \"params\": [\n",
    "                    p for n, p in param_optimizer if any(nd in n for nd in no_decay)\n",
    "                ],\n",
    "                \"weight_decay\": 0.0,\n",
    "            },\n",
    "        ]\n",
    "        optimizer = optim.AdamW(optimizer_parameters, lr=config.LR)\n",
    "\n",
    "    num_training_steps = len(train_dataloader) * config.EPOCHS\n",
    "    scheduler = get_linear_schedule_with_warmup(\n",
    "        optimizer,\n",
    "        num_warmup_steps=0,\n",
    "        num_training_steps=num_training_steps\n",
    "    )\n",
    "\n",
    "    min_avg_val_loss = float('inf')\n",
    "    for epoch in range(config.EPOCHS):\n",
    "        train(model, train_dataloader, valid_dataloader, criterion, optimizer, scheduler, epoch)\n",
    "        avg_val_loss = val(model, valid_dataloader, criterion)\n",
    "\n",
    "        if config.SAVE_BEST_ONLY:\n",
    "            if avg_val_loss < min_avg_val_loss:\n",
    "                best_model = copy.deepcopy(model)\n",
    "                best_val_mse_score = avg_val_loss\n",
    "\n",
    "                model_name = 'roberta_best_model'\n",
    "                torch.save(best_model.state_dict(), model_name + '.pt')\n",
    "\n",
    "                print(f'--- Best Model. Val loss: {min_avg_val_loss} -> {avg_val_loss}')\n",
    "                min_avg_val_loss = avg_val_loss\n",
    "\n",
    "    return best_model, best_val_mse_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T07:50:39.866435Z",
     "start_time": "2021-10-11T07:50:33.313134Z"
    },
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at hfl/chinese-roberta-wwm-ext were not used when initializing BertModel: ['cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.seq_relationship.bias']\n",
      "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Model(\n",
       "  (roberta): BertModel(\n",
       "    (embeddings): BertEmbeddings(\n",
       "      (word_embeddings): Embedding(21128, 768, padding_idx=1)\n",
       "      (position_embeddings): Embedding(512, 768)\n",
       "      (token_type_embeddings): Embedding(2, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.0, inplace=False)\n",
       "    )\n",
       "    (encoder): BertEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (1): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (2): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (3): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (4): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (5): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (6): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (7): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (8): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (9): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (10): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (11): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-07, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (pooler): BertPooler(\n",
       "      (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "      (activation): Tanh()\n",
       "    )\n",
       "  )\n",
       "  (dropout): Dropout(p=0.2, inplace=False)\n",
       "  (high_dropout): Dropout(p=0.5, inplace=False)\n",
       "  (attention): Sequential(\n",
       "    (0): Linear(in_features=768, out_features=768, bias=True)\n",
       "    (1): Tanh()\n",
       "    (2): Linear(in_features=768, out_features=1, bias=True)\n",
       "    (3): Softmax(dim=1)\n",
       "  )\n",
       "  (cls): Sequential(\n",
       "    (0): Linear(in_features=768, out_features=6, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = Model()\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T08:11:31.619230Z",
     "start_time": "2021-10-11T07:50:47.120753Z"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4349423d99694d0ab41d48e071ba12ac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Epoch 0', max=4138.0, style=ProgressStyle(description_wid…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-- Step: 1300\n",
      "Val loss: 0.46378527439158895\n",
      "-- Step: 2600\n",
      "Val loss: 0.391792626792322\n",
      "-- Step: 3900\n",
      "Val loss: 0.3687638635220735\n",
      "\n",
      "Training loss: 0.4100741486602097\n",
      "Val loss: 0.3676599715228962\n",
      "--- Best Model. Val loss: inf -> 0.3676599715228962\n"
     ]
    }
   ],
   "source": [
    "best_model, best_val_mse_score = run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T08:13:28.897967Z",
     "start_time": "2021-10-11T08:13:28.873858Z"
    }
   },
   "outputs": [],
   "source": [
    "dataset_size = len(test_df)\n",
    "test_indices = list(range(dataset_size))\n",
    "\n",
    "test_data = TransformerDataset(test_df, test_indices, set_type='test')\n",
    "test_dataloader = DataLoader(test_data, batch_size=config.BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T08:13:30.326868Z",
     "start_time": "2021-10-11T08:13:30.314613Z"
    }
   },
   "outputs": [],
   "source": [
    "def predict(model):\n",
    "    val_loss = 0\n",
    "    test_pred = []\n",
    "    model.eval()\n",
    "    for step, batch in enumerate(test_dataloader):\n",
    "        b_input_ids = batch['input_ids'].to(device)\n",
    "        b_attention_mask = batch['attention_mask'].to(device)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            logits = model(input_ids=b_input_ids, attention_mask=b_attention_mask)\n",
    "            logits = logits.cpu().numpy()\n",
    "            test_pred.extend(logits)\n",
    "\n",
    "    test_pred = np.array(test_pred)\n",
    "    return test_pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T08:17:04.263380Z",
     "start_time": "2021-10-11T08:13:31.742828Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>label_0</th>\n",
       "      <th>label_1</th>\n",
       "      <th>label_2</th>\n",
       "      <th>label_3</th>\n",
       "      <th>label_4</th>\n",
       "      <th>label_5</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>count</th>\n",
       "      <td>21376.000000</td>\n",
       "      <td>21376.000000</td>\n",
       "      <td>21376.000000</td>\n",
       "      <td>21376.000000</td>\n",
       "      <td>21376.000000</td>\n",
       "      <td>21376.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.035835</td>\n",
       "      <td>0.087174</td>\n",
       "      <td>0.106294</td>\n",
       "      <td>0.147099</td>\n",
       "      <td>0.114739</td>\n",
       "      <td>0.208310</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.122105</td>\n",
       "      <td>0.244001</td>\n",
       "      <td>0.263792</td>\n",
       "      <td>0.346998</td>\n",
       "      <td>0.223904</td>\n",
       "      <td>0.378798</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>min</th>\n",
       "      <td>-0.199498</td>\n",
       "      <td>-0.334330</td>\n",
       "      <td>-0.325735</td>\n",
       "      <td>-0.373501</td>\n",
       "      <td>-0.357324</td>\n",
       "      <td>-0.423927</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25%</th>\n",
       "      <td>-0.022906</td>\n",
       "      <td>-0.022440</td>\n",
       "      <td>0.005438</td>\n",
       "      <td>-0.010263</td>\n",
       "      <td>0.007866</td>\n",
       "      <td>0.012773</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50%</th>\n",
       "      <td>0.013354</td>\n",
       "      <td>0.038101</td>\n",
       "      <td>0.049858</td>\n",
       "      <td>0.048381</td>\n",
       "      <td>0.067941</td>\n",
       "      <td>0.101083</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>75%</th>\n",
       "      <td>0.060324</td>\n",
       "      <td>0.117724</td>\n",
       "      <td>0.111919</td>\n",
       "      <td>0.149329</td>\n",
       "      <td>0.144410</td>\n",
       "      <td>0.247564</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>max</th>\n",
       "      <td>2.628645</td>\n",
       "      <td>3.049479</td>\n",
       "      <td>3.595179</td>\n",
       "      <td>3.555378</td>\n",
       "      <td>3.077674</td>\n",
       "      <td>3.503212</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            label_0       label_1       label_2       label_3       label_4  \\\n",
       "count  21376.000000  21376.000000  21376.000000  21376.000000  21376.000000   \n",
       "mean       0.035835      0.087174      0.106294      0.147099      0.114739   \n",
       "std        0.122105      0.244001      0.263792      0.346998      0.223904   \n",
       "min       -0.199498     -0.334330     -0.325735     -0.373501     -0.357324   \n",
       "25%       -0.022906     -0.022440      0.005438     -0.010263      0.007866   \n",
       "50%        0.013354      0.038101      0.049858      0.048381      0.067941   \n",
       "75%        0.060324      0.117724      0.111919      0.149329      0.144410   \n",
       "max        2.628645      3.049479      3.595179      3.555378      3.077674   \n",
       "\n",
       "            label_5  \n",
       "count  21376.000000  \n",
       "mean       0.208310  \n",
       "std        0.378798  \n",
       "min       -0.423927  \n",
       "25%        0.012773  \n",
       "50%        0.101083  \n",
       "75%        0.247564  \n",
       "max        3.503212  "
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_pred = predict(best_model)\n",
    "\n",
    "tmp = pd.DataFrame(test_pred)\n",
    "tmp.columns = [f'label_{i}' for i in range(6)]\n",
    "tmp.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T08:17:50.252004Z",
     "start_time": "2021-10-11T08:17:50.243542Z"
    }
   },
   "outputs": [],
   "source": [
    "test_pred[test_pred < 0] = 0\n",
    "test_pred[test_pred > 3] = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T08:17:53.581359Z",
     "start_time": "2021-10-11T08:17:53.364682Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>emotion</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>34170_0002_A_12</td>\n",
       "      <td>0.0,0.0010360582,0.08128041,0.0,0.0081214905,0.0089266</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>34170_0002_A_14</td>\n",
       "      <td>0.0,0.0,0.0,0.0,0.0,0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>34170_0003_A_16</td>\n",
       "      <td>0.0,0.009026202,0.0,0.07286199,0.064760275,0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>34170_0003_A_17</td>\n",
       "      <td>0.0,0.0,0.0,0.023623193,0.0,0.22696266</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>34170_0003_A_18</td>\n",
       "      <td>0.0,0.04119603,0.0,0.0,0.11077476,0.15557021</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                id                                                 emotion\n",
       "0  34170_0002_A_12  0.0,0.0010360582,0.08128041,0.0,0.0081214905,0.0089266\n",
       "1  34170_0002_A_14                                 0.0,0.0,0.0,0.0,0.0,0.0\n",
       "2  34170_0003_A_16          0.0,0.009026202,0.0,0.07286199,0.064760275,0.0\n",
       "3  34170_0003_A_17                  0.0,0.0,0.0,0.023623193,0.0,0.22696266\n",
       "4  34170_0003_A_18            0.0,0.04119603,0.0,0.0,0.11077476,0.15557021"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test['emotion'] = list(test_pred)\n",
    "\n",
    "sub = submit.copy()\n",
    "del sub['emotion']\n",
    "\n",
    "sub = sub.merge(test[['id','emotion']], how='left', on='id')\n",
    "sub['emotion'] = sub['emotion'].apply(lambda x: ','.join([str(i) for i in x]))\n",
    "sub.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-11T08:18:03.955129Z",
     "start_time": "2021-10-11T08:18:03.877983Z"
    }
   },
   "outputs": [],
   "source": [
    "sub.to_csv(f'roberta_baseline.tsv', sep='\\t', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
